# candidate_generator.py
# -*- coding: utf-8 -*-

import logging
from typing import Dict, Any, List, Tuple, Set, Optional
from collections import Counter, defaultdict
import re # 用于 N-gram 提取
import os
import numpy as np # 用于处理向量数据
import faiss # 用于语义索引查询

# 导入项目模块
import config
from plagiarism_core.core_engine import ProcessedDocData # 导入预处理数据结构
from plagiarism_core.index_manager import IndexManager

# 获取日志记录器
logger = logging.getLogger(__name__)

class CandidateGenerator:
    """
    根据配置和可用的索引，为输入文档生成候选比对文档集合。
    """

    def __init__(self, index_manager: IndexManager):
        """
        初始化候选集生成器。

        Args:
            index_manager (IndexManager): 已初始化的索引管理器实例。
        """
        logger.info("初始化 CandidateGenerator...")
        self.index_manager = index_manager
        # 检查并记录可用索引的状态
        self.minhash_available = index_manager.minhash_enabled
        self.inverted_index_available = index_manager.inverted_index_enabled
        self.semantic_index_available = index_manager.semantic_index_enabled
        logger.info(f"候选集生成器可用策略: MinHash LSH={'是' if self.minhash_available else '否'}, "
                    f"倒排索引={'是' if self.inverted_index_available else '否'}, "
                    f"语义索引={'是' if self.semantic_index_available else '否'}")

    def generate(self, input_data: ProcessedDocData) -> Tuple[Set[str], Dict[str, int]]:
        """
        为给定的输入文档数据生成候选语料库文档 ID 集合。
        根据配置启用不同的召回策略，并合并结果。
        """
        task_id_str = f"[任务 {getattr(self.index_manager, 'current_task_id', '未知')}]" # 尝试获取任务ID
        logger.info(f"{task_id_str} 开始为输入文档 '{os.path.basename(input_data['source_file_path'])}' 生成候选集...")

        candidate_ids: Set[str] = set() # 存储最终的候选 ID
        recall_stats: Dict[str, int] = defaultdict(int) # 统计各策略召回数

        # --- 策略 1: MinHash LSH 召回 ---
        if config.USE_MINHASH_PRESCREENING:
            if not self.minhash_available:
                logger.warning(f"{task_id_str} 配置启用了 MinHash 但索引不可用，跳过 MinHash 召回。")
            else:
                minhash_candidates = self._get_candidates_minhash(input_data)
                count = len(minhash_candidates)
                candidate_ids.update(minhash_candidates)
                recall_stats["minhash_candidates"] = count
                logger.info(f"{task_id_str} MinHash LSH 召回了 {count} 个候选文档。")
        else:
             logger.debug(f"{task_id_str} 配置中禁用 MinHash LSH 召回。")

        # --- 策略 2: 倒排索引召回 ---
        if config.USE_INVERTED_INDEX:
            if not self.inverted_index_available:
                 logger.warning(f"{task_id_str} 配置启用了倒排索引但索引不可用，跳过倒排索引召回。")
            else:
                inverted_candidates = self._get_candidates_inverted_index(input_data)
                count = len(inverted_candidates)
                candidate_ids.update(inverted_candidates)
                recall_stats["inverted_index_candidates"] = count
                logger.info(f"{task_id_str} 倒排索引召回了 {count} 个候选文档。")
        else:
              logger.debug(f"{task_id_str} 配置中禁用倒排索引召回。")

        # --- 策略 3: 语义向量索引召回 ---
        if config.USE_SEMANTIC_INDEX:
            if not self.semantic_index_available:
                 logger.warning(f"{task_id_str} 配置启用了语义索引但索引不可用，跳过语义索引召回。")
            else:
                # 语义索引需要输入文档的块信息
                chunk_type = config.SEMANTIC_CHUNK_TYPE
                input_chunks = input_data.get("sentences_pos") if chunk_type == 'sentence' else input_data.get("paragraphs_pos", [])
                if not input_chunks:
                     logger.warning(f"{task_id_str} 未能在输入数据中找到用于语义索引查询的 '{chunk_type}' 块。")
                     semantic_candidates = set()
                else:
                    semantic_candidates = self._get_candidates_semantic_index(input_chunks)
                count = len(semantic_candidates)
                candidate_ids.update(semantic_candidates)
                recall_stats["semantic_index_candidates"] = count
                logger.info(f"{task_id_str} 语义索引召回了 {count} 个候选文档。")
        else:
               logger.debug(f"{task_id_str} 配置中禁用语义索引召回。")

        # --- 如果没有任何召回策略启用，则默认全选 ---
        if not config.USE_MINHASH_PRESCREENING and not config.USE_INVERTED_INDEX and not config.USE_SEMANTIC_INDEX:
             logger.warning(f"{task_id_str} 未启用任何召回策略！默认将所有语料库文档视为候选（这可能非常低效）。")
             all_corpus_ids = set(self.index_manager.corpus_data.keys())
             candidate_ids = all_corpus_ids
             recall_stats["no_strategy_all_selected"] = len(all_corpus_ids)

        logger.info(f"{task_id_str} 候选集生成完成。总候选文档数（去重后）: {len(candidate_ids)}")
        return candidate_ids, dict(recall_stats)


    # === 各召回策略的具体实现 ===

    def _get_candidates_minhash(self, input_data: ProcessedDocData) -> Set[str]:
        """使用 MinHash LSH 获取候选文档 ID。"""
        candidate_ids = set()
        task_id_str = f"[任务 {getattr(self.index_manager, 'current_task_id', '未知')}]"

        minhash_components = self.index_manager.get_minhash_lsh()
        if not minhash_components:
            logger.debug(f"{task_id_str} MinHash LSH 组件不可用，跳过查询。")
            return candidate_ids

        lsh_index, corpus_minhashes = minhash_components
        cleaned_input_text = input_data.get("cleaned_text", "")

        # 检查输入文本长度是否满足阈值
        if len(cleaned_input_text) < config.MINHASH_INPUT_MIN_CHARS_THRESHOLD:
             logger.info(f"{task_id_str} 输入文本长度 ({len(cleaned_input_text)}) 小于 MinHash 阈值 ({config.MINHASH_INPUT_MIN_CHARS_THRESHOLD})，跳过 MinHash 召回。")
             return candidate_ids

        # 计算输入文本的 MinHash (可以考虑从 input_data 中获取预计算的)
        input_minhash_sig = IndexManager._calculate_minhash(cleaned_input_text) # 调用静态方法
        if not input_minhash_sig:
            logger.warning(f"{task_id_str} 未能为输入文本计算有效的 MinHash 签名，跳过 MinHash LSH 查询。")
            return candidate_ids

        logger.debug(f"{task_id_str} 开始执行 MinHash LSH 查询...")
        try:
            # LSH query 返回的是 bytes 类型的 key
            result_keys_bytes = lsh_index.query(input_minhash_sig)
            # 解码为 utf-8 字符串
            minhash_candidates = {key.decode('utf-8', errors='ignore') for key in result_keys_bytes}

            # 过滤掉可能不在当前语料库中的 ID (如果 LSH 索引未及时更新)
            current_corpus_ids = set(self.index_manager.corpus_data.keys())
            valid_candidates = {cid for cid in minhash_candidates if cid in current_corpus_ids}
            invalid_count = len(minhash_candidates) - len(valid_candidates)
            if invalid_count > 0:
                logger.warning(f"{task_id_str} MinHash LSH 返回了 {invalid_count} 个不在当前语料库中的文档 ID（可能来自旧索引）。")

            candidate_ids.update(valid_candidates)
            logger.debug(f"{task_id_str} MinHash LSH 查询完成，找到 {len(valid_candidates)} 个有效候选。")
        except Exception as e:
            logger.exception(f"{task_id_str} 执行 MinHash LSH 查询时发生错误。")

        return candidate_ids


    def _get_candidates_inverted_index(self, input_data: ProcessedDocData) -> Set[str]:
        """使用倒排索引获取候选文档 ID。"""
        candidate_ids = set()
        task_id_str = f"[任务 {getattr(self.index_manager, 'current_task_id', '未知')}]"

        inverted_index = self.index_manager.get_inverted_index()
        if not inverted_index:
             logger.debug(f"{task_id_str} 倒排索引组件不可用，跳过查询。")
             return candidate_ids

        cleaned_input_text = input_data.get("cleaned_text", "")
        if not cleaned_input_text:
            logger.warning(f"{task_id_str} 输入文本为空，无法使用倒排索引查询。")
            return candidate_ids

        logger.debug(f"{task_id_str} 开始执行倒排索引查询...")
        try:
            # 1. 提取输入文本中的查询 N-grams
            query_ngrams = self._extract_query_ngrams(
                cleaned_input_text,
                config.INVERTED_INDEX_NGRAM_SIZE,
                config.INVERTED_INDEX_QUERY_NGRAM_COUNT
            )
            if not query_ngrams:
                logger.info(f"{task_id_str} 未能从输入文本中提取有效的查询 N-gram，跳过倒排索引查询。")
                return candidate_ids
            logger.debug(f"{task_id_str} 提取了 {len(query_ngrams)} 个查询 N-gram 用于倒排索引。")

            # 2. 查询索引并统计文档命中次数
            doc_match_counts = defaultdict(int)
            current_corpus_ids = set(self.index_manager.corpus_data.keys()) # 获取当前有效的语料ID
            for ngram in query_ngrams:
                # 从索引中获取包含该 N-gram 的文档列表 [(doc_id, pos), ...]
                postings = inverted_index.get(ngram, [])
                # 提取唯一的文档 ID
                # 同时过滤掉不存在于当前语料库的 doc_id (处理索引可能过时的情况)
                doc_ids_for_ngram = {doc_id for doc_id, _ in postings if doc_id in current_corpus_ids}
                # 增加每个命中 N-gram 的文档的计数
                for doc_id in doc_ids_for_ngram:
                    doc_match_counts[doc_id] += 1

            # 3. 根据阈值筛选候选文档
            candidate_threshold = config.INVERTED_INDEX_CANDIDATE_THRESHOLD
            matched_docs_inverted = {
                doc_id for doc_id, count in doc_match_counts.items()
                if count >= candidate_threshold
            }
            candidate_ids.update(matched_docs_inverted)
            logger.debug(f"{task_id_str} 倒排索引查询完成。共有 {len(doc_match_counts)} 个文档至少命中 1 个 N-gram，"
                         f"其中 {len(matched_docs_inverted)} 个文档命中次数 >= 阈值 {candidate_threshold}。")

        except Exception as e:
            logger.exception(f"{task_id_str} 执行倒排索引查询时发生错误。")

        return candidate_ids


    def _get_candidates_semantic_index(self, input_chunks_with_pos: List[Tuple[str, int, int]]) -> Set[str]:
        """使用语义向量索引 (Faiss) 获取候选文档 ID。"""
        semantic_candidate_docs = set()
        task_id_str = f"[任务 {getattr(self.index_manager, 'current_task_id', '未知')}]"

        semantic_components = self.index_manager.get_semantic_index()
        if not semantic_components:
             logger.debug(f"{task_id_str} 语义索引组件不可用，跳过查询。")
             return semantic_candidate_docs

        semantic_index, semantic_id_map = semantic_components

        if not input_chunks_with_pos:
            logger.warning(f"{task_id_str} 输入块列表为空，无法执行语义索引查询。")
            return semantic_candidate_docs

        logger.debug(f"{task_id_str} 开始执行语义向量索引查询 (输入块数: {len(input_chunks_with_pos)})...")
        try:
            # --- 1. 获取输入块的嵌入向量 ---
            # 需要 SemanticComparer 实例来编码输入块
            semantic_comparer = getattr(self.index_manager, 'semantic_comparer_instance', None)
            if not semantic_comparer or not semantic_comparer.enabled:
                 logger.error(f"{task_id_str} 无法获取有效的 SemanticComparer 实例来编码输入块，语义索引查询失败。")
                 return semantic_candidate_docs

            input_chunk_texts = [item[0] for item in input_chunks_with_pos]
            logger.debug(f"  正在编码 {len(input_chunk_texts)} 个输入块用于语义查询...")
            # 编码时确保使用与构建索引时相同的模型和设备
            input_embeddings = semantic_comparer.model.encode(
                input_chunk_texts,
                convert_to_numpy=True, # Faiss 需要 numpy array
                show_progress_bar=False,
                device=semantic_comparer.device,
                batch_size=config.SEMANTIC_ENCODE_BATCH_SIZE
            )
            input_embeddings_np = np.array(input_embeddings).astype('float32')

            # --- L2 归一化查询向量 (如果索引向量已归一化) ---
            # （构建索引时已归一化，查询时也需要归一化以保证内积=余弦相似度）
            faiss.normalize_L2(input_embeddings_np)
            logger.debug(f"  输入块向量已编码并归一化 (形状: {input_embeddings_np.shape})。")

            # --- 2. 执行 Faiss 搜索 ---
            k_neighbors = config.SEMANTIC_INDEX_TOP_K
            logger.debug(f"  开始在 Faiss 索引 (向量数: {semantic_index.ntotal}) 中搜索 Top {k_neighbors} 邻居...")
            # search 方法返回 (距离/得分, 索引ID)
            # 对于内积（已归一化），distances 是余弦相似度得分
            distances, indices = semantic_index.search(input_embeddings_np, k_neighbors)
            logger.debug(f"  Faiss 搜索完成。返回距离形状: {distances.shape}, 索引形状: {indices.shape}")

            # --- 3. 处理搜索结果，筛选并映射回文档 ID ---
            similarity_threshold = config.SEMANTIC_INDEX_CANDIDATE_THRESHOLD
            logger.debug(f"  开始处理搜索结果，筛选得分 >= {similarity_threshold:.4f} 的匹配...")
            candidate_vector_indices = set() # 存储满足条件的向量在索引中的 ID
            total_hits_above_threshold = 0

            # 遍历每个输入块的搜索结果
            for i in range(indices.shape[0]): # i 是输入块的索引
                for j in range(indices.shape[1]): # j 是邻居的索引 (0 to k-1)
                    vector_idx = indices[i, j]  # 获取邻居向量在 Faiss 索引中的 ID (0 到 ntotal-1)
                    score = distances[i, j]    # 获取对应的得分 (余弦相似度)

                    # Faiss 返回 -1 表示没有更多邻居或该位置无效
                    if vector_idx == -1:
                        break # 当前查询块的后续邻居无效

                    # 应用相似度阈值
                    if score >= similarity_threshold:
                        candidate_vector_indices.add(vector_idx)
                        total_hits_above_threshold += 1
                    else:
                        # 由于邻居按得分降序排列，后续得分只会更低，提前中断
                        break

            logger.debug(f"  结果筛选完成。共找到 {total_hits_above_threshold} 个向量匹配（来自 {len(candidate_vector_indices)} 个唯一向量索引）。")

            # --- 将向量索引映射回文档 ID ---
            current_corpus_ids = set(self.index_manager.corpus_data.keys())
            retrieved_doc_count = 0
            for vector_idx in candidate_vector_indices:
                if 0 <= vector_idx < len(semantic_id_map):
                    try:
                        # 从映射中获取文档 ID
                        map_info = semantic_id_map[vector_idx]
                        corpus_id = map_info[0]
                        # 确保该文档 ID 仍然存在于当前加载的语料库中
                        if corpus_id in current_corpus_ids:
                            if corpus_id not in semantic_candidate_docs:
                                 semantic_candidate_docs.add(corpus_id)
                                 retrieved_doc_count += 1
                        # else: # 可选：记录找到但已不存在的文档
                        #     logger.debug(f"语义索引找到向量 {vector_idx} 对应已移除/无效的文档 ID '{corpus_id}'")
                    except IndexError:
                        logger.warning(f"语义索引映射错误：无法访问 semantic_id_map[{vector_idx}] (映射大小: {len(semantic_id_map)})")
                    except Exception as map_err:
                         logger.error(f"处理语义索引映射信息时出错 (索引: {vector_idx}, 映射信息: {map_info}): {map_err}")
                else:
                    logger.warning(f"语义索引返回了越界的向量索引: {vector_idx} (有效范围 0-{len(semantic_id_map)-1})")

            logger.debug(f"{task_id_str} 语义索引查询完成。共召回 {len(semantic_candidate_docs)} 个唯一候选文档。")

        except faiss.FaissException as faiss_err:
            logger.exception(f"{task_id_str} 执行 Faiss 搜索过程中发生错误: {faiss_err}")
        except AttributeError as attr_err:
             logger.exception(f"{task_id_str} 语义索引查询时发生属性错误 (可能 semantic_comparer 未正确初始化?): {attr_err}")
        except Exception as e:
            logger.exception(f"{task_id_str} 执行语义索引查询时发生未知错误。")

        return semantic_candidate_docs


    # === 辅助方法 ===

    def _extract_query_ngrams(self, text: str, ngram_size: int, count: int) -> List[str]:
        """
        从输入文本中提取用于倒排索引查询的 N-grams。
        策略：提取所有 N-grams，计算频率，选择出现频率适中的前 count 个。
        （此策略可根据效果调整，例如优先选择 TF-IDF 值高的 N-gram）
        """
        if not text or len(text) < ngram_size or count <= 0:
            return []

        logger.debug(f"开始从输入文本中提取最多 {count} 个查询 N-gram (大小={ngram_size})...")
        all_ngrams = []
        # 生成所有可能的 N-grams
        for i in range(len(text) - ngram_size + 1):
            ngram = text[i:i + ngram_size]
            # 过滤掉包含换行符或完全是空白的 ngram
            if '\n' not in ngram and ngram.strip():
                # 可以添加更复杂的过滤规则，例如避免全是标点
                # if re.search(r'\w', ngram, re.UNICODE): # 检查是否包含单词字符
                 all_ngrams.append(ngram)

        if not all_ngrams:
            logger.warning("未能从输入文本中提取任何有效的 N-gram 用于查询。")
            return []

        # 计算 N-gram 频率
        ngram_counts = Counter(all_ngrams)
        total_unique_ngrams = len(ngram_counts)
        logger.debug(f"共找到 {len(all_ngrams)} 个 N-gram 实例，其中 {total_unique_ngrams} 个是唯一的。")

        # --- 选择策略：优先选择中等频率的 N-gram ---
        # 设定频率阈值（可调整）
        min_freq_threshold = 2  # 至少出现两次
        # 最大频率阈值：例如，不超过总 N-gram 实例数的 5% 或至少 5 次
        max_freq_threshold = max(5, int(len(all_ngrams) * 0.05))

        # 筛选出频率在阈值范围内的 N-gram
        filtered_ngrams = [
            ng for ng, freq in ngram_counts.items()
            if min_freq_threshold <= freq <= max_freq_threshold
        ]
        logger.debug(f"筛选出 {len(filtered_ngrams)} 个中等频率 N-gram (频率范围: [{min_freq_threshold}, {max_freq_threshold}])。")

        # 如果筛选后数量不足，可以放宽条件或补充最高频的
        if len(filtered_ngrams) < count:
            logger.debug(f"中频 N-gram 数量 ({len(filtered_ngrams)}) 不足 {count}，将补充其他 N-gram...")
            # 按频率降序排序所有唯一 N-gram
            sorted_by_freq = sorted(ngram_counts.items(), key=lambda item: item[1], reverse=True)
            # 从高频开始补充，直到达到 count 或用尽所有 N-gram
            additional_ngrams = [ng for ng, freq in sorted_by_freq if ng not in filtered_ngrams]
            filtered_ngrams.extend(additional_ngrams)

        # 最后截取所需数量
        selected_ngrams = filtered_ngrams[:count]

        logger.info(f"最终选择了 {len(selected_ngrams)} 个 N-gram 用于倒排索引查询。")
        # logger.debug(f"选中的查询 N-grams: {selected_ngrams}") # 调试时可以打印
        return selected_ngrams
